package mpt

import (
	"sync"

	"github.com/ethereum/go-ethereum/common"
	lru "github.com/hashicorp/golang-lru"
	"hundsun.com/hsl/hschain/common/config"
	"hundsun.com/hsl/hschain/common/log"
	"hundsun.com/hsl/hschain/common/types"
	pbcom "hundsun.com/hsl/hschain/protos/common"
	dbm "hundsun.com/hsl/hschain/store/db"
	"hundsun.com/hsl/hschain/store/state"
)

const (
	// DefaultCacheSize 默认缓存mpt树
	DefaultCacheSize = 16
	// StateName ...
	StateName = "mpt"
)

// Store mpt store backend
type Store struct {
	db    dbm.Database
	wraps *sync.Map // map[string]*TrieWrap
	cache *lru.Cache
}

func init() {
	state.Register(StateName, New)
}

// New new mpt store
func New(cfg *config.PluginConfig, db dbm.Database) (state.State, error) {
	cache, _ := lru.New(DefaultCacheSize)
	s := &Store{db: db, wraps: &sync.Map{}, cache: cache}
	return s, nil
}

// Get get values by keys
func (s *Store) Get(datas *pbcom.StateGet) ([][]byte, error) {
	var wrap *TrieWrap
	var err error
	values := make([][]byte, len(datas.Keys))
	strHash := string(datas.StateHash)
	if data, ok := s.cache.Get(strHash); ok {
		wrap = data.(*TrieWrap)
	} else if data, ok := s.wraps.Load(strHash); ok {
		wrap = data.(*TrieWrap)
	} else {
		wrap, err = NewTrieWrap(s.db, datas.StateHash)
		if nil != err {
			log.Errorf("store Get NewTrieWrap err %v, root hash %s", err, common.Bytes2Hex(datas.StateHash))
			return values, err
		}
		s.cache.Add(strHash, wrap)
	}
	if err == nil {
		for i := 0; i < len(datas.Keys); i++ {
			value, err := wrap.Get(datas.Keys[i])
			if nil == err {
				values[i] = value
			}
		}
	}
	return values, nil
}

// GetStateInfo 获取最新状态数据信息
func (s *Store) GetStateInfo() (*pbcom.StateInfo, error) {
	infoDB := state.NewInfoDB(s.db)
	return infoDB.GetStateInfo()
}

// SetMemory set keys values to memory mpt, return root hash and error
func (s *Store) SetMemory(datas *pbcom.StateSet, sync bool) ([]byte, error) {
	wrap, err := NewTrieWrap(s.db, datas.PreStateHash)
	if nil != err {
		log.Errorf("store SetMemory NewTrieWrap err %v, root hash %s", err, common.Bytes2Hex(datas.PreStateHash))
		return nil, err
	}

	for i := 0; i < len(datas.KVs); i++ {
		wrap.Update(datas.KVs[i].Key, datas.KVs[i].Value)
	}

	hash, err := wrap.Commit()
	if err != nil {
		log.Errorf("store SetMemory Commit to memory trie fail %v", err)
		return nil, err
	}
	wrap.SetStateInfo(&pbcom.StateInfo{
		Ty:        StateName,
		Number:    datas.Number,
		StateHash: hash,
	})
	s.wraps.Store(string(hash), wrap)
	return hash, nil
}

// Commit convert memory mpt to storage db
func (s *Store) Commit(hash []byte) error {
	v, ok := s.wraps.Load(string(hash))
	if !ok {
		log.Errorf("store mpt commit error %v", types.ErrNotFound)
		return types.ErrNotFound
	}
	wrap := v.(*TrieWrap)
	err := wrap.Commit2DB(hash)
	if nil != err {
		log.Errorf("store mpt commit error %v", err)
		return err
	}
	// append LatestStateInfoKV
	info := wrap.GetStateInfo()
	kv, err := state.InfoKV(info)
	if err != nil {
		log.Errorf("store mpt commit LatestStateInfoKV error %v", err)
		return err
	}
	s.db.Set(kv.Key, kv.Value)
	s.wraps.Delete(string(hash))
	return nil
}

// RollBack rollback data and sign, mpt only need rollback state info
func (s *Store) RollBack(datas *pbcom.StateSet) error {
	if datas == nil {
		return nil
	}
	var preNumber uint64
	if datas.Number > 0 {
		preNumber = datas.Number - 1
	}
	info := &pbcom.StateInfo{
		Ty:        StateName,
		Number:    preNumber,
		StateHash: datas.PreStateHash,
	}
	kv, err := state.InfoKV(info)
	if err != nil {
		log.Errorf("store mpt rollback LatestStateInfoKV error %v", err)
		return err
	}
	s.db.Set(kv.Key, kv.Value)
	return nil
}
